[2/4] Diffusion Quantized ckpt export #810
Conversation
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
jingyu-ml
changed the title
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #810 +/- ##
==========================================
+ Coverage 73.31% 73.40% +0.08%
==========================================
Files 192 193 +1
Lines 19613 19911 +298
==========================================
+ Hits 14380 14616 +236
- Misses 5233 5295 +62 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
| else: | ||
| cpu_state_dict = { | ||
| k: v.detach().contiguous().cpu() for k, v in component.state_dict().items() | ||
| } | ||
| save_file(cpu_state_dict, str(component_export_dir / "model.safetensors")) | ||
| with open(component_export_dir / "config.json", "w") as f: | ||
| json.dump( | ||
| { | ||
| "_class_name": type(component).__name__, | ||
| "_export_format": "safetensors_state_dict", | ||
| }, | ||
| f, | ||
| indent=4, | ||
| ) |
There was a problem hiding this comment.
Can we combine these with L851 to L863? They look duplicated.
Why we need to offload tensors to cpu before saving?
There was a problem hiding this comment.
if we always save with safetensors, keeping the .cpu() is the safe/default choice. this is also how the transformers/diffusers save_pretrained save the tensors to safetensors file.
There was a problem hiding this comment.
Could you clarify more?
Can we combine these with L851 to L863? They look duplicated.
Line 880 saves the state dict to safe tensor, line 884 saves the quant config to config.json. we use these 2 function only if the model is not diffusers based.
There was a problem hiding this comment.
cpu_state_dict = {
k: v.detach().contiguous().cpu() for k, v in component.state_dict().items()
}
save_file(cpu_state_dict, str(component_export_dir / "model.safetensors"))
with open(component_export_dir / "config.json", "w") as f:
json.dump(
{
"_class_name": type(component).__name__,
"_export_format": "safetensors_state_dict",
},
f,
indent=4,
)
I mean this code block appears twice in the same script.
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
jingyu-ml
force-pushed
the
jingyux/2-3-diffusion-export
branch
from
ef4f814 to
9f0e998
Compare
ChenhanYu
left a comment
ChenhanYu
left a comment
There was a problem hiding this comment.
Commented on the dynamic module part.
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
LGTM, left a few more comments.
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
4 participants
What does this PR do?
Type of change: New feature
Overview:
This MR adds HuggingFace checkpoint export support for LTX‑2 by treating TI2VidTwoStagesPipeline as a diffusion-like pipeline, exporting only the stage‑1 transformer (with QKV-fusion-enabled dummy inputs) and falling back to writing model.safetensors when save_pretrained isn’t available. It also preserves the original forward in DynamicModule patching (_forward_pre_dm) so downstream callers can still invoke the pre-patched forward implementation.
Changes
DynamicModulepatching: when patching forward, we now stash the pre-patched implementation inself._forward_pre_dm(once) so downstream code can still call the original forward, then re-bind forward to the class implementation. This is needed for the LTX2 FP8 calibration.export_hf_checkpoint()now also treats ltx_pipelines.ti2vid_two_stages.TI2VidTwoStagesPipeline as a “diffusion-like” object and routes it through _export_diffusers_checkpoint() (import guarded; no hard dependency).Plans
Usage
Testing
Before your PR is "Ready for review"
Additional Information
Summary by CodeRabbit
Release Notes
New Features
--extra-paramCLI option for flexible model configuration and parameter passingChores
✏️ Tip: You can customize this high-level summary in your review settings.